# -*- coding: utf-8 -*-

import threading
from app.foundation import pro_api
from app.enum.stock import ListStatusEnum, ExchangeEnum, MarketEnum
from app.models import db, Stock, TradeCalendar, StockName, Index, StockIncome
from app.decorator import freq_limit, command_log


@command_log
def get_stock_list():
    """
    获取股票列表
    :return: void
    """
    data = pro_api.stock_basic(list_status=ListStatusEnum.LIST, fields=','.join(
        ['ts_code', 'symbol', 'name', 'area', 'industry', 'fullname', 'enname', 'market', 'exchange', 'curr_type',
         'list_status', 'list_date', 'delist_date', 'is_hs']))
    # 替换None字段
    data['area'].fillna(value='', inplace=True)
    data['industry'].fillna(value='', inplace=True)
    data['is_hs'].fillna(value='', inplace=True)
    stocks = []
    for index, item in data.iterrows():
        stock = Stock.update_or_create({'ts_code': item['ts_code']}, item)
        stocks.append(stock)
    db.session.bulk_save_objects(stocks)
    db.session.commit()


@command_log
def get_index_list():
    """
    获取指数列表
    """
    for market in MarketEnum.TEXT.keys():
        data = pro_api.index_basic(market=market, fields=','.join(
            ['ts_code', 'name', 'fullname', 'market', 'publisher', 'index_type', 'category', 'base_date', 'base_point',
             'list_date', 'weight_rule', 'desc', 'exp_date']))
        # 替换None字段
        data['fullname'].fillna(value='', inplace=True)
        data['market'].fillna(value='', inplace=True)
        data['publisher'].fillna(value='', inplace=True)
        data['index_type'].fillna(value='', inplace=True)
        data['category'].fillna(value='', inplace=True)
        data['base_date'].fillna(value='', inplace=True)
        data['base_point'].fillna(value=0.00, inplace=True)
        data['list_date'].fillna(value='', inplace=True)
        data['weight_rule'].fillna(value='', inplace=True)
        data['exp_date'].fillna(value='', inplace=True)
        indexes = []
        for i, item in data.iterrows():
            index = Index.update_or_create({'ts_code': item['ts_code']}, item)
            indexes.append(index)
        db.session.bulk_save_objects(indexes)
        db.session.commit()


@command_log
def get_trade_calendar():
    """
    获取交易日历
    """
    # 多线程拉取数据
    thread_list = []
    for exchange in ExchangeEnum.TEXT.keys():
        t = threading.Thread(target=get_exchange_trade_calendar, args=(exchange,))
        thread_list.append(t)
    for t in thread_list:
        t.setDaemon(True)
        t.start()
    for t in thread_list:
        t.join()


def get_exchange_trade_calendar(exchange):
    data = pro_api.trade_cal(exchange=exchange, fields=','.join(['exchange', 'cal_date', 'is_open', 'pretrade_date']))
    # 替换None字段
    data['pretrade_date'].fillna(value='', inplace=True)
    calendars = []
    for index, item in data.iterrows():
        calendar = TradeCalendar.update_or_create({'exchange': item['exchange'], 'cal_date': item['cal_date']},
                                                  item)
        calendars.append(calendar)
    db.session.bulk_save_objects(calendars)
    db.session.commit()


@command_log
def get_name_change():
    """
    获取股票曾用名
    """
    stocks = Stock.query.with_entities(Stock.ts_code).all()
    for stock in stocks:
        get_stock_name_change(ts_code=stock[0])


@freq_limit(2)
def get_stock_name_change(ts_code):
    data = pro_api.namechange(ts_code=ts_code, fields=','.join(
        ['ts_code', 'name', 'start_date', 'end_date', 'ann_date', 'change_reason']))
    if data.empty:
        return
    # 替换None字段
    data['end_date'].fillna(value='', inplace=True)
    data['ann_date'].fillna(value='', inplace=True)
    data['change_reason'].fillna(value='', inplace=True)
    stock_names = []
    for index, item in data.iterrows():
        stock_name = StockName.update_or_create({'ts_code': item['ts_code'], 'start_date': item['start_date']}, item)
        stock_names.append(stock_name)
    db.session.bulk_save_objects(stock_names)
    db.session.commit()


@command_log
def get_income():
    stocks = Stock.query.with_entities(Stock.ts_code).all()
    for stock in stocks:
        get_stock_income(stock[0])


@freq_limit(1.3)
def get_stock_income(ts_code):
    """
    获取股票利润接口数据，每分钟最多请求50次
    """
    data = pro_api.income(ts_code=ts_code, fields=','.join(
        ['ts_code', 'ann_date', 'f_ann_date', 'end_date', 'report_type', 'comp_type', 'total_revenue', 'revenue',
         'total_cogs', 'oper_cost', 'operate_profit',
         'total_profit', 'n_income', 'update_flag']))
    # 替换None字段
    data['ann_date'].fillna(value='', inplace=True)
    data['f_ann_date'].fillna(value='', inplace=True)
    data['report_type'].fillna(value=0, inplace=True)
    data['comp_type'].fillna(value=0, inplace=True)
    data['total_revenue'].fillna(value=0.0000, inplace=True)
    data['revenue'].fillna(value=0.0000, inplace=True)
    data['total_cogs'].fillna(value=0.0000, inplace=True)
    data['oper_cost'].fillna(value=0.0000, inplace=True)
    data['operate_profit'].fillna(value=0.0000, inplace=True)
    data['total_profit'].fillna(value=0.0000, inplace=True)
    data['n_income'].fillna(value=0.0000, inplace=True)
    data['update_flag'].fillna(value=0, inplace=True)
    stock_incomes = []
    for index, item in data.iterrows():
        stock_income = StockIncome.update_or_create({'ts_code': item['ts_code'], 'end_date': item['end_date']}, item)
        stock_incomes.append(stock_income)
    db.session.bulk_save_objects(stock_incomes)
    db.session.commit()
